from pymilvus import MilvusClient
import asyncio
from typing_extensions import override
from .BasePool import BasePool, Entity, PoolConfig


class MilvusPool(BasePool[MilvusClient]):
    def __init__(
        self,
        uri: str,
        db_name: str = "default",
        username: str = None,
        password: str = None,
        config: PoolConfig = None,
    ):
        super().__init__(config)
        self.uri = uri
        self.db_name = db_name
        self.username = username
        self.password = password

    @override
    async def create_connection(self):
        if len(self._pool) >= self.config.max_connections:
            self.config.logger.warning(
                f"已达到最大连接数 ({self.config.max_connections}), 无法创建新连接"
            )
            return None

        # 重试逻辑
        for attempt in range(self.config.connection_retry + 1):
            try:
                # 创建连接
                client = MilvusClient(
                    uri=self.uri,
                    user=self.username,
                    password=self.password,
                    db_name=self.db_name,
                )

                connection = Entity[MilvusClient](client=client)

                self.config.logger.debug(
                    f"成功创建新的Milvus连接 (当前连接数: {len(self._pool) + 1})"
                )
                if await self.validate_connection(connection):
                    return connection

            except Exception as e:
                if attempt < self.config.connection_retry:
                    self.config.logger.warning(
                        f"连接Milvus失败 (尝试 {attempt + 1}/{self.config.connection_retry + 1}): {str(e)}"
                    )
                    await asyncio.sleep(self.config.retry_interval)
                else:
                    self.config.logger.error(
                        f"连接Milvus失败，已达到最大重试次数: {str(e)}"
                    )
                    return None

    
    async def validate_connection(self, connection: Entity[MilvusClient]) -> bool:
        # TODO 添加验证逻辑
        return True
